# python -m lilab.OpenLabCluster_train.t2_halfmirror_hierachy *.clippredpkl
# %%
import pickle
import numpy as np
import matplotlib.pyplot as plt
import os
import argparse

def np_norm(array):
    return np.linalg.norm(array, axis=-1)

clippredpkl = '/DATA/taoxianming/rat/data/Mix_analysis/SexAgeDay55andzzcWTinAUT_MMFF/result32/semiseq2seq_iter0/output/semisupervise-decSeq-iter2-epoch5/olc-2024-05-23-semiseq2seq-iter2-epoch5_pca12.clippredpkl'

def main(clippredpkl):
    project = os.path.dirname(clippredpkl)
    out_dir = os.path.join(project, 'halfmirror_hierachy')
    os.makedirs(out_dir, exist_ok=True)

    clippreddata = pickle.load(open(clippredpkl, 'rb'))

    assert clippreddata.keys() > set(['nK_mutual', 'nK_mirror_half', 'embedding', 'cluster_labels', 'cluster_names_mutualmerge'])

    # %%
    # 寻找每个 cluster 的embedding 中心点
    embedding = clippreddata['embedding']
    embedding_d2 = clippreddata['embedding_d2']
    cluster_labels = clippreddata['cluster_labels'] - 1 #start from 0
    nK = clippreddata['ncluster']
    nK_mutual = clippreddata['nK_mutual']
    nK_mirror_half = clippreddata['nK_mirror_half']
    cluster_names_mutualmerge = np.array(clippreddata['cluster_names_mutualmerge'])
    # 寻找每个 cluster 的embedding 中心点
    embedding_center = np.zeros((nK, embedding.shape[1]))
    for i in range(nK):
        embedding_center[i] = np.mean(embedding[cluster_labels == i], axis=0)

    embedding_center_mutual = embedding_center[:nK_mutual]
    embedding_center_mirror_half = embedding_center[nK_mutual:nK_mutual+nK_mirror_half]
    embedding_center_mirror_mutual_dist = np_norm(embedding_center_mirror_half[:,None] - embedding_center_mutual[None,:])
    embedding_center_mirror_mutual_dist_avg = np.mean(embedding_center_mirror_mutual_dist, axis=-1)
    ind_mirror_half_closest_to_mutual = np.argsort(embedding_center_mirror_mutual_dist_avg)

    print(cluster_names_mutualmerge[nK_mutual:][ind_mirror_half_closest_to_mutual])

    seed_mutual = ind_mirror_half_closest_to_mutual[-2] + nK_mutual
    seed_pos = embedding_center_mirror_half[seed_mutual - nK_mutual]
    embedding_center_mirror = np.array([embedding_center[nK_mutual:nK_mutual+nK_mirror_half],
                                        embedding_center[nK_mutual+nK_mirror_half:]])
    embedding_center_2mirror = embedding_center[nK_mutual:]

    if False:
        #%% 方法1
        embedding_disttoseed_mirror = np_norm(embedding_center_mirror - seed_pos[None,None])
        embedding_indtoseed_mirror = np.argmin(embedding_disttoseed_mirror, axis=0)

        ind_mirror_half = np.take_along_axis(np.arange(nK_mutual, nK).reshape(2,nK_mirror_half), embedding_indtoseed_mirror[None], axis=0).ravel()
    else:
        #%% 方法2
        recruit_l = [seed_mutual - nK_mutual]
        candidate_l = np.arange(nK_mirror_half)
        while len(candidate_l) > 0:
            embedding_center_seeds = embedding_center_2mirror[recruit_l] #nrecruit, ndim
            embedding_center_candidate = embedding_center_mirror[:, candidate_l] #2, ncandidate, ndim
            dist = np_norm(embedding_center_seeds[:,None,None] - embedding_center_candidate[None,]) #nrecruit, 2, ncandidate
            idx_min = np.argmin(dist.ravel())
            _, id_lr, id_candidate = np.unravel_index(idx_min, dist.shape)
            recruit_l.append(candidate_l[id_candidate] + nK_mirror_half*id_lr)
            candidate_l = np.delete(candidate_l, id_candidate)

        ind_mirror_half = np.array(sorted(recruit_l[1:], key=lambda x: x if x<nK_mirror_half else x-nK_mirror_half)) + nK_mutual

    #%%
    ind_mutualmirror_half = np.concatenate([np.arange(nK_mutual), ind_mirror_half])
    d2_x0, d2_y0 = embedding_d2.T

    embeddingd2_center = np.zeros((nK, embedding_d2.shape[1]))
    for i in range(nK):
        embeddingd2_center[i] = np.mean(embedding_d2[cluster_labels == i], axis=0)


    plt.figure(figsize=(4,4))
    plt.plot(d2_x0[::10], d2_y0[::10], '.', color='gray')

    embedding_center_mutualmirror_half = embeddingd2_center[ind_mutualmirror_half]
    d2_x, d2_y = embedding_center_mutualmirror_half.T
    plt.plot(d2_x[:nK_mutual], d2_y[:nK_mutual], 'ro')
    plt.plot(d2_x[nK_mutual:], d2_y[nK_mutual:], 'bo')
    plt.plot(d2_x[seed_mutual], d2_y[seed_mutual], 'go')

    for i in range(nK_mirror_half):
        plt.plot(*embeddingd2_center[[nK_mutual+i, nK_mutual+nK_mirror_half+i]].T, 'k-')

    plt.savefig(os.path.join(out_dir, 'embedding_d2.png'))


    #%%
    #save clippredpkl
    ind_valid = np.isin(cluster_labels, ind_mutualmirror_half)

    c = clippreddata['cluster_labels'].copy()
    c[c > nK_mutual+nK_mirror_half] -= nK_mirror_half
    clippreddata_new = {'embedding': clippreddata['embedding'][ind_valid],
                        'embedding_d2': clippreddata['embedding_d2'][ind_valid],
                        'cluster_labels': c[ind_valid],
                        'cluster_names': cluster_names_mutualmerge,
                        'ncluster': nK_mutual + nK_mirror_half,
                        'ntwin': clippreddata['ntwin'],
                        'clipNames': clippreddata['clipNames'][ind_valid],
                        'df_clipNames': clippreddata['df_clipNames'].loc[ind_valid].reset_index(drop=True)}
    out_clippredfile = os.path.join(out_dir, 'clippreddata_halfmirror.clippredpkl')
    pickle.dump(clippreddata_new, open(out_clippredfile, 'wb'))
    print('save to ', out_dir)
    os.system(f'python -m lilab.openlabcluster_postprocess.s3b_hiecluster_plot_pro "{out_clippredfile}"')


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('clippredpkl', type=str)
    args =  parser.parse_args()
    main(args.clippredpkl)
